import numpy as np
import torch
from scipy.spatial.distance import directed_hausdorff
from torch import nn


def dice_coefficient_3D(predict, target, num_classes, epsilon=1e-6):
    # The first output of this network has been reshaped to (-1,class_num). Softmax operation is also performed

    if isinstance(predict, torch.Tensor) and isinstance(target, torch.Tensor):
        total_dice = 0
        for i in range(num_classes):
            # if len(target.shape) > 2:
            #     target = target[:, i, :, :, :].reshape(-1, 1)
            # if len(predict.shape) > 2:
            #     predict = predict[:, i, :, :, :].reshape(-1, 1)
            intersection = torch.sum(torch.mul(predict, target))
            debug_predict_sum = torch.sum(predict)
            debug_target_sum = torch.sum(target)
            dice = (2 * intersection + epsilon) / (debug_predict_sum + debug_target_sum + epsilon)
            total_dice += dice
        return total_dice / num_classes
    elif isinstance(predict, np.ndarray) and isinstance(target, np.ndarray):
        total_dice = 0
        for i in num_classes:
            target = target[:, i, :]
            predict = predict[:, i, :]
            intersection = np.sum(np.multiply(predict, target))
            debug_predict_sum = np.sum(predict)
            debug_target_sum = np.sum(target)
            dice = (2 * intersection + epsilon) / (debug_predict_sum + debug_target_sum + epsilon)
            total_dice += dice
        return total_dice / num_classes
    else:
        raise TypeError


def Hausdorff_Distance_3D(predict, target):
    assert (predict.shape == target.shape)
    if isinstance(predict, torch.Tensor):
        predict = predict.clone().detach().cpu().numpy()
    if isinstance(target, torch.Tensor):
        target = target.clone().detach().cpu().numpy()
    Hausdorff_Distance: float = 0.
    compute_num = 0
    if len(predict.shape) == 5:
        for batch_num in range(predict.shape[0]):
            for channel_num in range(predict.shape[1]):
                for z_index in range(predict.shape[4]):
                    one_layer_HD = directed_hausdorff(predict[batch_num, channel_num, :, :, z_index],
                                                      target[batch_num, channel_num, :, :, z_index])[0]
                    compute_num += 1
                    Hausdorff_Distance += one_layer_HD
        return Hausdorff_Distance, round(Hausdorff_Distance / compute_num, 2)

    else:
        for z_index in range(predict.shape[2]):
            one_layer_HD = directed_hausdorff(predict[:, :, z_index],
                                              target[:, :, z_index])[0]
            compute_num += 1
            Hausdorff_Distance += one_layer_HD
        return Hausdorff_Distance, round(Hausdorff_Distance / compute_num, 2)


class Dice(nn.Module):
    def __init__(self, class_num, smooth=1e-6, threshold=None, weights=None):
        super(Dice, self).__init__()
        if class_num < 1:
            raise Exception("Wrong class number")
        self.class_num = class_num
        self.smooth = smooth
        self.weights = weights
        self.threshold = threshold

    def _binary_dice(self, predict, target):
        if self.threshold is not None:
            predict = torch.where(predict > self.threshold, torch.tensor(1).cuda(), torch.tensor(0).cuda())
        predict = predict.contiguous().view(-1)
        target = target.contiguous().view(-1)
        intersection = torch.sum((predict * target))
        result = ((2. * intersection + self.smooth) / (torch.sum(predict) + torch.sum(target) + self.smooth))
        return result

    def forward(self, predict, target):
        target = target.long()
        if self.class_num == 1:
            return self._binary_dice(predict, target)
        else:
            if target.shape[1] == 1:
                processed_target = torch.zeros([target.shape[0], self.class_num] + list(target.shape[2:]))
                for i in range(target.shape[0]):  # Considering that the batch size may be greater than 1
                    tmp_label = torch.nn.functional.one_hot(target[i], num_classes=self.class_num)
                    processed_target[i] = tmp_label[0].permute(3, 0, 1, 2)
            else:
                processed_target = target
            channel_num = processed_target.shape[1]
            # if weights is None:
            # 	weights = torch.ones(C) #uniform weights for all classes
            total_dice = 0
            dice_list = []
            for i in range(1, channel_num):
                dice = self._binary_dice(predict[:, i, :, :, :], processed_target[:, i, :, :, :])
                dice_list.append(dice)
                # print("dice",i,"：",diceLoss.sum()/len(diceLoss))
                if self.weights is not None:
                    dice *= self.weights[i]
                total_dice += dice
            return total_dice / (channel_num - 1), dice_list


def DSC(predict, target, threshold):
    epsilon = 1e-6
    # In the KiPA dataset 0 is the background
    assert (predict.shape == target.shape)
    if isinstance(predict, torch.Tensor) and isinstance(target, torch.Tensor):
        predict = torch.where(predict > threshold, 1, 0)
        TP = torch.sum(torch.multiply(predict, target)).item()
        FN = torch.sum(torch.multiply(1 - predict, target)).item()
        FP = torch.sum(torch.multiply(predict, 1 - target)).item()
        dsc = (2 * TP + epsilon) / (2 * TP + FN + FP + epsilon)
        if dsc > 1.0:
            print("A")
        return dsc
    elif isinstance(predict, np.ndarray) and isinstance(target, np.ndarray):
        predict = np.where(predict > threshold, 1, 0)
        TP = np.sum(np.multiply(predict, target))
        FN = np.sum(np.multiply(1 - predict, target))
        FP = np.sum(np.multiply(predict, 1 - target))
        dsc = (2 * TP + epsilon) / (2 * TP + FN + FP + epsilon)
        return dsc
    else:
        raise TypeError


def get_wrong_voxels(predict, target):
    if isinstance(predict, torch.Tensor) and isinstance(target, torch.Tensor):
        predict = predict.clone().detach().cpu().numpy()
        target = target.clone().detach().cpu().numpy()
    predict_correct = np.multiply(predict, target)
    predict_omit = np.multiply(1 - predict, target)
    predict_redundant = np.multiply(predict, 1 - target)
    return predict_correct, predict_omit, predict_redundant
